# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于双量子点量子模型
单量子比特的归零操作
多个初始态对多个目标态的归零保真度

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
import copy

np.random.seed(1)
T = np.pi
dt = np.pi/20
step_max = T/dt
sx = np.mat([[0, 1], [1, 0]], dtype=complex)
sy = np.mat([[0, -1j], [1j, 0]], dtype=complex)
sz = np.mat([[1, 0], [0, -1]], dtype=complex)

# a0,a1,a2,a3,a4,a5,a6,a7,a8,a9 = 0,0,0,0,0,0,0,0,0,0 #统计各动作被选用的频率

action_space = np.mat([[1,0,0], #可以选择的动作范围，各列的每项分别代表着 sigma x, y, z 前面的系数。
                       [2,0,0], #每次执行的动作都是单独的绕 x, y, z 轴一定角度的旋转
                       [0,1,0], # x, y 方向的值可以取负，但 z 方向的只能取正值
                       [0,2,0],
                       [0,0,1],
                       [0,0,2],
                       [-1,0,0],
                       [-2,0,0],
                       [0,-1,0],
                       [0,-2,0]])

theta_num = 6 #除了 0 和 Pi 两个点之外，点的数量
varphi_num = 21#varphi 角度一圈上的点数

theta = np.linspace(0,np.pi,theta_num+2,endpoint=True) 
varphi = np.linspace(0,np.pi*2,varphi_num,endpoint=False) 

def psi_set():
    psi_set = []
    for ii in range(1,theta_num+1):
        for jj in range(varphi_num):
            psi_set.append(np.mat([[np.cos(theta[ii]/2)],[np.sin(theta[ii]/2)*(np.cos(varphi[jj])+np.sin(varphi[jj])*(0+1j))]]))
    psi_set.append(np.mat([[1], [0]], dtype=complex))
    psi_set.append(np.mat([[0], [1]], dtype=complex))
    return psi_set

target_set = psi_set()
init_set = psi_set()

#动作直接选最优的
def step0(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        psi_list_ = copy.deepcopy(psi_list)
        fid_list_ = copy.deepcopy(fid_list)
        
        del psi_list_[fid_list_.index(max(fid_list_))]
        del fid_list_[fid_list_.index(max(fid_list_))]
        
        best_action = fid_list.index(max(fid_list_))
        
        best_fid = max(fid_list_)
        
        # best_action = fid_list.index(min(fid_list))
        # best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

def job(target_psi):
    fids_list = []
    for psi1 in init_set:
        
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        
        fid_max = F
        fid_max1 = F
        fid_max2 = F
        fid_max0 = F
        
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step0(psi,target_psi,F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max1,fid_max2,fid_max0)  
        fids_list.append(fid_max)
        
    return  np.mean(fids_list)

def multicore():
    pool = mp.Pool()
    F_list = pool.map(job, target_set)
    return F_list
    

if __name__ == '__main__':
    # print(target_set)
    time1 = time()
    F_list = multicore()
    print(F_list)
    time2 = time()
    print(np.mean(F_list))
    print('time_cost is: ',time2-time1)

# 128测试点
#动作 x,y: 0,1,2,-1,-2; z: 0,1,2

# [0.9992626016568953, 0.999514484401934, 0.9995460533336445, 0.9994487703895798, 0.9994341030357903, 0.9994277361138151, 0.9995065968535353, 0.9995340631758372, 0.9995306864977319, 0.999426204281705, 0.9993963640685807, 0.999411472525704, 0.9995366303277334, 0.9993891174005957, 0.9994904903874451, 0.9994201956923201, 0.9993321384538574, 0.9995540107497178, 0.9994607633414887, 0.9995125756924348, 0.9993758571776079, 0.9992989411366318, 0.9994344492569684, 0.9995358246179075, 0.9995188853013948, 0.9994997090832396, 0.9993183356366728, 0.9994400424435175, 0.9994868779795223, 0.999538155729168, 0.9995312501317175, 0.9993789657781853, 0.9994442339539455, 0.9994548054567861, 0.999477453460133, 0.9994847719313813, 0.9994484890569231, 0.9994418498151304, 0.9995079291848901, 0.9994711079783859, 0.9995770571775666, 0.999485094422466, 0.9994551053468359, 0.9995230980837229, 0.9994664195841043, 0.9994094370995563, 0.9995110420974335, 0.9994362014931573, 0.9995731027656826, 0.9995489943208492, 0.9990514218806035, 0.9995147454366133, 0.9994823498684318, 0.9994760814110144, 0.999577758784197, 0.9990079323385279, 0.9994297317319802, 0.9994418324776184, 0.9994075732558478, 0.9995401661775851, 0.9992999654172481, 0.9995206790708213, 0.9995324703488049, 0.9994551053468359, 0.9995230980837231, 0.9994664195841046, 0.9994094370995564, 0.9995110420974336, 0.9994362014931574, 0.9995731027656828, 0.9995489943208493, 0.9990492450563018, 0.9995147454366134, 0.9994823498684318, 0.9994760814110146, 0.9995777587841973, 0.999007932338528, 0.9994297317319805, 0.9994418324776186, 0.9994075732558481, 0.9995401661775851, 0.9992999654172482, 0.9995206790708215, 0.999532470348805, 0.9993013807849245, 0.9994344492569686, 0.9995358246179076, 0.9995188853013948, 0.9994997090832395, 0.9993183356366728, 0.9994400424435175, 0.9994868779795223, 0.999538155729168, 0.9995312501317175, 0.9993789657781853, 0.9994442339539457, 0.9994548054567862, 0.999477453460133, 0.9994847719313812, 0.9994484890569231, 0.9994418498151303, 0.9995079291848901, 0.999471107978386, 0.9995770571775666, 0.999485094422466, 0.9992626016568953, 0.9995144844019338, 0.9995460533336447, 0.9994487703895798, 0.9994341030357902, 0.9994277361138151, 0.9995065968535355, 0.9995340631758372, 0.999530686497732, 0.999426204281705, 0.9993963640685806, 0.999411472525704, 0.9995366303277334, 0.9993891174005957, 0.999490490387445, 0.9994201956923203, 0.9993321384538573, 0.9995540107497178, 0.9994607633414886, 0.9995125756924349, 0.9993758571776077, 0.998921714734468, 0.9989193984437936]
# 0.9994439365950352
# time_cost is:  80.03747224807739


